# -*- coding: utf-8 -*-
"""
Created on Sun Jan 16 09:39:37 2022

@author: 懒得很
@revisor: Sage_713705
"""
import pandas as pd
import numpy as np
import os
import matplotlib
import matplotlib.pyplot as plt
import seaborn as sns
import sys

def getmaxaffinity(fpath):
    try:
        with open(fpath, encoding="utf-8") as f:
            for line in f:
                if line.strip().startswith("REMARK VINA RESULT:"):
                    parts = line.strip().split()
                    if len(parts) >= 4:
                        return parts[3]  # parts[3] 是亲和力数值
        return ''
    except Exception as e:
        print(f"读取文件出错: {fpath}, 错误: {e}")
        return ''

if __name__ == "__main__":
    # 提供数据来源路径
    f = input("请输入待处理对接数据的存放路径：").strip()
    if not os.path.isdir(f):
        print("路径不存在，请检查！")
        sys.exit(1)

    # 操作 pdbqt 文件
    receptors = []
    ligands = []
    affinity = []
    for root,dirs,files in os.walk(f):
        for file in files:
            if os.path.splitext(file)[1]== '.pdbqt':  # 修改操作文件类型为 .pdbqt 而非 .txt
                fn = os.path.splitext(file)[0]
                try:
                    receptor,ligand = fn.split('_2_')
                    receptors.append(receptor)
                    ligands.append(ligand)
                    fpath = os.path.join(f,file)
                    affinity.append(getmaxaffinity(fpath))
                except:
                    continue
                
    # 过滤掉不能转为 float 的亲和力数据
    valid_receptors = []
    valid_ligands = []
    aff_n = []
    for r, l, a in zip(receptors, ligands, affinity):
        # print(f"原始affinity内容: {a}")  # 调试，查看亲和力原始字符串
        try:
            aff_n.append(float(a))
            valid_receptors.append(r)
            valid_ligands.append(l)
        except Exception:
            print(f"{r}与{l}无对接结果")
            continue

    df = pd.DataFrame()
    df['receptor'] = valid_receptors
    df['ligand'] = valid_ligands
    df['affinity'] = aff_n

    # 调试用代码，打印数据长度测试是否为空数据
    # print("receptors:", len(receptors))
    # print("ligands:", len(ligands))
    # print("affinity:", len(aff_n))
    # print(df.head())

    # 先转为字符串类型再分割
    df['ligand'] = df['ligand'].astype(str).str.split('_', expand=True).iloc[:, 0]
    # 无需进行亲和力数据筛选
    # df = df.loc[df['affinity'] < -7.0]
    df = df.pivot(index = 'ligand', columns = 'receptor', values = 'affinity')
    # 空数据判断
    if df.empty:
        print("没有有效的数据，无法绘制热图！")
        input("按回车键退出……")
        sys.exit(1)

    # 绘制美化后的热图
    plt.figure(figsize=(12, 6))  # 设置图片大小
    ax = sns.heatmap(
        df,
        cmap='OrRd_r',                # 选择配色方案
        annot=True,                   # 显示每个格子的数值
        fmt=".1f",                    # 数值保留一位小数
        linewidths=0.5,               # 单元格间隔线宽
        linecolor='white',            # 单元格间隔线颜色
        cbar_kws={'label': 'Binding Affinity (kcal/mol)'},  # 色带标签
        annot_kws={"size": 12}        # 数值字体大小
    )
    ax.set_xlabel("receptor", fontsize=14)  # x轴标签及字体
    ax.set_ylabel("ligand", fontsize=14)    # y轴标签及字体
    ax.set_xticklabels(ax.get_xticklabels(), rotation=45, ha='right', fontsize=12)  # x轴刻度美化
    ax.set_yticklabels(ax.get_yticklabels(), rotation=0, fontsize=12)              # y轴刻度美化
    plt.tight_layout()  # 自动调整布局防止标签重叠

    # 保存数据到根目录下的 result 文件夹
    root_dir = os.path.dirname(os.path.abspath(__file__))  # 获取脚本所在目录
    save_dir = os.path.join(root_dir, "result")            # 拼接 result 文件夹路径
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)                              # 若不存在则创建
    save_path = os.path.join(save_dir, "heatmap.png")      # 拼接图片完整路径
    plt.savefig(save_path, dpi=300)                        # 保存图片，分辨率300dpi
    print(f"热图已保存到: {save_path}")
    plt.close()                                            # 关闭绘图窗口

    # 全局报错检索驻留
    input("按回车键退出……")